import numpy as np


def distributed_auction_with_virtual_usv(H, epsilon=10):
    H = np.asarray(H, dtype=float)
    m, n = H.shape  # m=USV行, n=UAV列

    if m >= n:
        # 真实USV足够，直接调用原算法 # TODO 会爆炸！需要检查！
        return distributed_auction(H, epsilon)

    size = n  # 扩到 n 行，保证每架UAV都有一个“行”可选
    H_pad = np.zeros((size, size), dtype=float)
    H_pad[:m, :n] = H  # 下面的虚拟行默认0效益
    H_pad[m:, :] = 1.5 * np.max(H)
    # H_pad[m:, :] = 11451.4

    assignment_pad, iters = distributed_auction(H_pad, epsilon)
    assignment_pad = np.asarray(assignment_pad, dtype=int)  # 长度 n，1-index

    # 标记虚拟USV：行号>m 的置为114514
    assignment = np.where(assignment_pad <= m, assignment_pad, 114514)
    return assignment, iters


def distributed_auction(H, epsilon=0.05):
    # Get the dimensions of the cost matrix H
    Cost_matrix = H
    m, n = Cost_matrix.shape

    # 邻居拓扑为全连通
    Neighbours = {i: [j for j in range(n) if j != i] for i in range(n)}

    # Initialize variables
    alpha = np.zeros(n, dtype=int)
    price_list = np.zeros((m, n))
    bidder_max = np.zeros((m, n), dtype=int)

    # Initialize price list and assignments
    for i in range(n):
        price = np.zeros(m)
        b = np.zeros(m, dtype=int)

        # Sort and determine the price for the object
        ST = np.sort(H[:, i])[::-1]
        price_object = ST[0] - ST[1]
        argmax = np.argmax(H[:, i])
        price[argmax] = price_object + np.random.rand() * 0.001
        price_list[:, i] = price

        b[argmax] = i + 1
        bidder_max[:, i] = b

        alpha[i] = argmax + 1

    iteration = 0
    price_list_previous = np.zeros((m, n))

    while True:
        iteration += n

        for agent in range(n):
            price_list_0 = price_list[:, agent]

            # Compute beta (100 - H[:, agent] as in MATLAB)
            beta = 100 * np.ones(m) - H[:, agent]

            # Get neighboring agents' price lists
            L = Neighbours[agent]
            price_list_t = price_list[:, L]

            alpha_0 = alpha[agent] - 1  # Convert 1-index to 0-index
            bidd = bidder_max[:, agent]
            bidd_neighbour = bidder_max[:, L]

            # Perform auction step
            price_list_update, alpha_t, bidd_update = Auction(agent, price_list_0, price_list_t, alpha_0, beta, bidd,
                                                              bidd_neighbour, epsilon, m)

            price_list[:, agent] = price_list_update
            alpha[agent] = alpha_t + 1  # Convert 0-index back to 1-index
            bidder_max[:, agent] = bidd_update

        if np.array_equal(price_list_previous, price_list):
            break

        price_list_previous = price_list.copy()

    assignment = alpha[alpha != 0]
    # TODO 与agents的id兼容, 在gc.py里面写这个转换函数
    return assignment, iteration


def Auction(agent, price_list_0, price_list_t, alpha_0, beta, bidd, bidd_neighbour, epsilon, n):
    price_list_update = np.zeros(n)
    bidd_max = np.zeros(n, dtype=int)

    for j in range(n):
        combined_prices = np.hstack((price_list_0[j], price_list_t[j]))
        argvalue = np.max(combined_prices)
        argmax = np.argmax(combined_prices)

        price_list_update[j] = argvalue

        bidd_acc = np.hstack((bidd[j], bidd_neighbour[j]))
        bidd_max[j] = bidd_acc[argmax]

    if price_list_0[alpha_0] <= price_list_update[alpha_0] and bidd_max[alpha_0] != agent + 1:
        argmax = np.argmax(beta - price_list_update)
        alpha_t = argmax
        bidd_max[argmax] = agent + 1  # Convert 0-index to 1-index

        bidd_update = bidd_max

        price_sort = np.sort(beta - price_list_update)[::-1]
        gamma = price_sort[0] - price_sort[1]
        price_list_update[alpha_t] = price_list_0[alpha_t] + gamma + epsilon
    else:
        alpha_t = alpha_0
        bidd_update = bidd_max

    return price_list_update, alpha_t, bidd_update



def verify_minimization(H, epsilon=0.05, atol=1e-6):
    """验证分布式拍卖结果是否达到(在当前实现语义下的)最小化目标。

    - 当 m < n 时，用与算法一致的方式在行方向补虚拟USV(成本=virtual_cost)
    - 与 Hungarian(或回退的全排列) 的全局最优做对比
    - 返回包含我们算法成本、最优成本、是否匹配、两组分配等信息
    """
    H = np.asarray(H, dtype=float)
    m, n = H.shape

    assignment, iters = distributed_auction_with_virtual_usv(H, epsilon)

    # 构造与算法一致的代价矩阵（最小化）
    if m < n:
        size = n
        H_pad = np.zeros((size, size), dtype=float)
        H_pad[:m, :n] = H
        H_pad[m:, :] = 2 * np.max(H)
    else:
        H_pad = H.copy()

    # 计算我们算法的总成本：对虚拟分配用第一行虚拟成本（所有虚拟行成本相同）
    total_cost_ours = 0.0
    for j, a in enumerate(assignment):
        if a == 114514:
            row_index = m if m < H_pad.shape[0] else H_pad.shape[0] - 1
            total_cost_ours += 0
        else:
            total_cost_ours += float(H_pad[int(a) - 1, j])

    # 计算全局最优（Hungarian，若不可用则小规模全排列回退）
    try:
        from scipy.optimize import linear_sum_assignment  # type: ignore
        row_ind, col_ind = linear_sum_assignment(H_pad)
        total_cost_opt = float(H_pad[row_ind, col_ind].sum())
        assignment_opt = np.full(n, 114514, dtype=int)
        for r, c in zip(row_ind, col_ind):
            if c < n:
                assignment_opt[c] = (int(r) + 1) if r < m else 114514
    except Exception:
        import itertools
        size_cols = H_pad.shape[1]
        best_cost = float("inf")
        best_perm = None
        # 仅适用于规模较小情形
        for perm in itertools.permutations(range(H_pad.shape[0]), size_cols):
            cost = 0.0
            for j in range(size_cols):
                cost += float(H_pad[perm[j], j])
            if cost < best_cost:
                best_cost = cost
                best_perm = perm
        total_cost_opt = float(best_cost)
        assignment_opt = np.array([(p + 1 if p < m else 114514) for p in best_perm], dtype=int)

    ok = bool(np.isclose(total_cost_ours, total_cost_opt, atol=atol))
    return {
        "ours_cost": total_cost_ours,
        "optimal_cost": total_cost_opt,
        "cost_match": ok,
        "ours_assignment": assignment.tolist(),
        "optimal_assignment": assignment_opt.tolist(),
        "iterations": int(iters),
    }

if __name__ == "__main__":
    np.random.seed(0)
    # 行 < 列 的简单测试用例（USV数量 < UAV数量）
    H = np.array([
        [80, 60, 55, 30],
        [70, 65, 40, 20],
    ], dtype=float)

    assignment, iterations = distributed_auction_with_virtual_usv(H, epsilon=0.05)

    # print("H:\n", H)
    # print("assignment (每架UAV对应的USV，1-index；114514表示分配到虚拟USV):", assignment)
    # print("iterations:", iterations)

    # mapping = [
    #     (uav_index + 1, (None if usv_id == 114514 else int(usv_id)))
    #     for uav_index, usv_id in enumerate(assignment)
    # ]
    # print("UAV->USV 映射(USV为None表示虚拟):", mapping)

    # # 验证最小化效果
    # result = verify_minimization(H, epsilon=0.05)
    # print("验证结果:", result)